#include "util/radix_tree/radix_tree.h"
#include "base/common.h"

namespace hawking {
namespace indexlib {

RadixTreeNode::RadixTreeNode(uint8_t power_of_slotnum, uint8_t height) :
    slice_array_(nullptr), power_of_slicenum_(power_of_slotnum),
    height_(height), shift_(power_of_slicenum_ * height_), mask_((1 << power_of_slicenum_) - 1) {}

void RadixTreeNode::Init(Pool* pool, RadixTreeNode* first_node) {
    uint32_t slice_num = SliceNum();
    slice_array_ = (void**)pool->Allocate(slice_num * sizeof(void*));
    slice_array_[0] = first_node;
    for (uint32_t idx = 1; idx < slice_num; ++idx) {
        slice_array_[idx] = nullptr;
    }
}

void RadixTreeNode::Insert(void* slot, uint64_t slice_id, Pool* pool) {
    uint32_t slice_index = ExtractSubSlotId(slice_id);
    if (!height_) {
        slice_array_[slice_index] = slot;
        return;
    }
        
    RadixTreeNode*& current_slot = (RadixTreeNode*&)slice_array_[slice_index];
    if (!current_slot) {
        void* buffer = pool->Allocate(sizeof(RadixTreeNode));
        RadixTreeNode* temp_node = (RadixTreeNode*)new (buffer) RadixTreeNode(power_of_slicenum_, height_ - 1);
        temp_node->Init(pool, nullptr);
        current_slot = temp_node;
    }

    current_slot->Insert(slot, slice_id, pool);
}

RadixTree::RadixTree(
    uint32_t slice_num_per_node, uint32_t slice_bytes, Pool* pool, uint16_t item_size) :
        pool_(pool), slice_bytes_(slice_bytes), item_size_(item_size) {
    DoInit_(slice_num_per_node, slice_bytes);
}

void RadixTree::DoInit_(uint32_t slice_num_per_node, uint32_t slice_bytes) {
    power_of_slicenum_ = CaclPowerOf2_(slice_num_per_node);
    slice_cursor_ = slice_bytes_;
    AllocateConsecutiveSlices_(1);
}

RadixTreeNode* RadixTree::CreateNode_(uint8_t height, RadixTreeNode* first_sub_node) {
    void* buffer = pool_->Allocate(sizeof(RadixTreeNode));
    RadixTreeNode* node = (RadixTreeNode*)new (buffer) RadixTreeNode(power_of_slicenum_, height);
    node->Init(pool_, first_sub_node);
    return node;
}

void RadixTree::AppendSlice_(void* slice) {
    if (!root_) {
        root_ = CreateNode_(0, nullptr);
    }

    uint8_t height = root_->Height();
    if (NeedGrowUp_(height)) {
        RadixTreeNode* node = CreateNode_(height + 1, root_);
        root_ = node;
    }

    root_->Insert(slice, slice_num_, pool_);
    ++slice_num_;
}

uint8_t* RadixTree::AllocateConsecutiveSlices_(uint32_t slice_num) {
    if (slice_num <= 0) {
        return nullptr;
    }

    uint8_t* current_slice = (uint8_t*)pool_->Allocate(slice_num * slice_bytes_);
    if (!current_slice) {
        return nullptr;
    }

    slice_cursor_ = 0;
    uint8_t* tmp_slice = current_slice;
    for (uint32_t idx = 0; idx < slice_num; ++idx) {
        AppendSlice_((void*)tmp_slice);
        tmp_slice += slice_bytes_;
    }

    return current_slice;
}

uint8_t* RadixTree::Allocate(size_t append_size) {
    uint8_t* current_slice = nullptr;

    if (append_size > slice_bytes_ - slice_cursor_) {
        uint32_t slice_num_to_allocate = CalculateNeededSliceNum_(append_size);
        current_slice = AllocateConsecutiveSlices_(slice_num_to_allocate);
        if (!current_slice) {
            LOG(ERROR) << "RadixTree allocate data fail to alloc " << slice_num_to_allocate;
            return nullptr;
        }

        slice_cursor_ = append_size % slice_bytes_;
    } else if (likely(root_ != nullptr)) {
        current_slice = (uint8_t*)root_->Search(slice_num_ - 1) + slice_cursor_;
        if (!current_slice) {
            LOG(ERROR) << "RadixTree Search fail";
            return nullptr;
        }
        slice_cursor_ += append_size;
    }

    return current_slice;
}

bool RadixTree::Append(const uint8_t* data, size_t length) {
    if (unlikely(!length || !data)) {
        return false;
    }

    uint8_t* address = Allocate(length);
    if (likely(address != nullptr)) {
        std::memcpy(address, data, length);
        return true;
    }

    return false;
}

}
}
